% 基于casadi的内点法ipopt；RK4离散等式约束; 返回值 矩阵 [v_ms x Ts Tr Fe Fm]
function res = atc_sqp_opt_v3(tvxfhg_ref, weights, LineInfo, TrainInfo, QsInfo)
pos_st_m = QsInfo.pos_start * 1000;
pos_ed_m = QsInfo.pos_end * 1000;
tvxfhg = tvxfhg_ref;
T_limit = 145;
abc = TrainInfo.abc;
m = TrainInfo.mass;
mt_num = TrainInfo.motor_num;
y = TrainInfo.rf;
dt = QsInfo.step_t;
grd_t = tvxfhg(:,6); %gradient
h_t = tvxfhg(:,5); %height
g = 9.81;
%% 系数准备
% 热路参数
Rs = 25.21994135;      % 定子热阻
Cs = 122.77;      % 定子热容
R_sr = 23.998;    % 气隙热阻
Rr = 87.781;      % 转子热阻
Cr = 89.345;       % 转子热容
% 等效定子损耗 Ps = k1 * Fe^2 + k3 * v^2 等价k1 k3系数，其中Fe,v的单位是 kN , m/s
% compute_coeff计算得
k1 = 0.076243065150734;
k3 = 0.001827167411627;
% 等效转子损耗 Pr = k1 * Fe^2 + k3 * v^2 等价k2 k4系数，其中Fe,v的单位是 kN , m/s
k2 = 0.251602114997421;
k4 = 2.030186012918816e-04;
% 等效风扇对定子等效冷却功率系数 Pws = d1 + k5*h + k6*(Ts-Tenv) cftool拟合得, 单位是W
% 所以变成kW要系数除以1000 负功率 
d1 = -1136.4 / 1000;
k5 = 1.4101 / 1000;
k6 = -78.2121 / 1000;
% 等效风扇对转子等效冷却功率系数 Pwr = d2 + k7*h + k8*(Tr-Tenv)
d2 = -1.5634e+03 / 1000;
k7 = 1.6876 / 1000;
k8 = -75.1086 / 1000;
% 状态方程系数
cs1 = -1/Cs*(1/Rs+1/R_sr-k6);
cs2 = 1/(Cs*R_sr);
cs3 = 1/Cs*k3;
cs4 = 1/Cs*k1;
cs5 = 1/Cs*(1/Rs-k6);

cr1 = 1/(Cr*R_sr);
cr2 = -1/Cr*(1/Rr+1/R_sr-k8);
cr3 = 1/Cr*k4;
cr4 = 1/Cr*k2;
cr5 = 1/Cr*(1/Rr-k8);

params = struct(...
    'abc', abc, ...          % 空气阻力系数 [a, b, c]（单位需与模型匹配）
    'm', m, ...              % 列车质量 (kg)
    'g', 9.81, ...           % 重力加速度 (m/s²)，若未定义变量g可直接赋值
    'y', y, ...              % 旋转质量系数（惯性系数）
    'cs1', cs1, ...          % 定子温度方程：定子温度耦合项系数
    'cs2', cs2, ...          % 定子温度方程：转子温度耦合项系数
    'cs3', cs3, ...          % 定子温度方程：速度平方项热生成系数
    'cs4', cs4, ...          % 定子温度方程：控制量U平方项损耗系数
    'cs5', cs5, ...          % 定子温度方程：环境温度耦合系数
    'Cs', Cs, ...            % 定子热容（单位：J/°C）
    'd1', d1, ...            % 定子散热系数（单位：W/°C）
    'k5', k5, ...            % 定子温度方程：高度相关散热修正系数
    'cr1', cr1, ...          % 转子温度方程：转子温度耦合项系数
    'cr2', cr2, ...          % 转子温度方程：定子温度耦合项系数
    'cr3', cr3, ...          % 转子温度方程：速度平方项热生成系数
    'cr4', cr4, ...          % 转子温度方程：控制量U平方项损耗系数
    'cr5', cr5, ...          % 转子温度方程：环境温度耦合系数
    'Cr', Cr, ...            % 转子热容（单位：J/°C）
    'd2', d2, ...            % 转子散热系数（单位：W/°C）
    'k7', k7, ...            % 转子温度方程：高度相关散热修正系数
    'mt_num', mt_num, ...    % 电机数量（用于归一化控制量U的平方项）
    'T_amb', 20 ...          % 环境温度（原代码中硬编码的20，建议参数化）
);
%% casadi
s_t_start = tic;
%长度
N = length(tvxfhg);
opti = casadi.Opti();
%定义状态变量 v s Ts Tr 4*N 单位m/s m 摄氏度
X = opti.variable(4,N);
%定义控制量 Fe Fm 2*N 单位kN
U = opti.variable(2,N);
%等式约束 矢量构造
A = diag(-ones(N, 1), 0) + diag(ones(N-1, 1), 1); %带状对角阵
A(end,:) = [];
B = abs(A);

e1 = abc(1)*g/1000/(1+y);
e2 = abc(2)*3.6*g/1000/(1+y);
e3 = abc(3)*3.6^2*g/1000/(1+y);
% %牛二等式
% opti.subject_to( A*X(1,:)' == ((U(1,1:N-1)'-U(2,1:N-1)')./(m*(1+y)) - e1.*ones(N-1,1) - e2.*X(1,1:N-1)' - e3.*X(1,1:N-1)'.^2 - ...
%     grd_t(1:N-1)*g/1000/(1+y)).*dt ); 
%位移积分线性等价
% opti.subject_to( A*X(2,:)' == 1/2.*B*X(1,:)'.*dt ); 
% %温升等效离散线性化 定子
% opti.subject_to( A*X(3,:)' == (cs1*X(3,1:N-1)' + cs2*X(4,1:N-1)' + cs3*X(1,1:N-1)'.^2 + cs4/mt_num^2.*U(1,1:N-1)'.^2 + ...
%     cs5*(20.*ones(N-1,1)-6/1000.*h_t(1:N-1)) - 1/Cs*(d1.*ones(N-1,1)+k5.*h_t(1:N-1)) ).*dt ); 
% %温升等效离散线性化 转子
% opti.subject_to( A*X(4,:)' == (cr1*X(3,1:N-1)' + cr2*X(4,1:N-1)' + cr3*X(1,1:N-1)'.^2 + cr4/mt_num^2.*U(1,1:N-1)'.^2 + ...
%     cr5*(20.*ones(N-1,1)-6/1000.*h_t(1:N-1)) - 1/Cr*(d2.*ones(N-1,1)+k7.*h_t(1:N-1)) ).*dt );

for i=2:N
    X_prev = X(:,i-1);  % 当前状态
    U_prev = U(:,i-1);  % 当前控制量
    h_t_prev = h_t(i-1);
    grd_t_prev = grd_t(i-1);
    dt = QsInfo.step_t;

    % --- 计算RK4斜率 ---
    % k1
    k1 = ode_system(X_prev, U_prev, h_t_prev, grd_t_prev, params) * dt;
    
    % k2（中间状态1）
    X_k2 = X_prev + 0.5*k1;
    k2 = ode_system(X_k2, U_prev, h_t_prev, grd_t_prev, params) * dt;
    
    % k3（中间状态2）
    X_k3 = X_prev + 0.5*k2;
    k3 = ode_system(X_k3, U_prev, h_t_prev, grd_t_prev, params) * dt;
    
    % k4（中间状态3）
    X_k4 = X_prev + k3;
    k4 = ode_system(X_k4, U_prev, h_t_prev, grd_t_prev, params) * dt;
    
    % --- RK4更新公式 ---
    X_next = X_prev + (k1 + 2*k2 + 2*k3 + k4)/6;
    
    % --- 添加约束 ---
    opti.subject_to(X(:,i) == X_next);
end
%不等式约束 矢量构造
%限速
opti.subject_to( 0 <= X(1,:) <= 200/3.6 );
%温度约束
opti.subject_to( X(3,:) < T_limit );
opti.subject_to( X(4,:) < T_limit );
%线性化电牵约束
opti.subject_to( U(1,:) <= -0.375*X(1,:)*3.6+276 );
opti.subject_to( U(1,:) <= -0.6341*X(1,:)*3.6+305.6 );
%线性化电制约束
opti.subject_to( U(1,:) >= -27.2*X(1,:)*3.6 );
opti.subject_to( U(1,:) >= -(-0.08667*X(1,:)*3.6+272.9) );
opti.subject_to( U(1,:) >= -(-0.7331*X(1,:)*3.6+358.8) );
%空气制动约束常数
opti.subject_to( 0 <= U(2,:) <= 300 );
%终端约束 起点约束
opti.subject_to( X(:,1) == [QsInfo.v_start/3.6; pos_st_m; QsInfo.T_st(1); QsInfo.T_st(2)] );
%终点约束
opti.subject_to( X(1,N) == QsInfo.v_end/3.6 ); %停车速度约束
opti.subject_to( pos_ed_m-50 <= X(2,N) <= pos_ed_m+50 ); %位置约束

s_t_end = toc(s_t_start);
fprintf("-约束构造耗时：%f s\n",s_t_end);

%% 目标函数 J
c_f_start = tic;
w1 = weights(1);
w2 = weights(2);
w3 = weights(3);
w4 = weights(4);
w5 = weights(5);
w6 = weights(6);
%目标函数：J = w1*(v-v_set)^2 + w2*(s-s_set)^2 + w3*(theta_s^2+theta_r^2) + U'*R*U 
J = 0;
for i=1:N
    % J = J + w1*(X(1,i)-tvxfhg(i,2)/3.6)^2 + w2*(X(2,i)-tvxfhg(i,3))^2 + w3*(X(3,i)+X(4,i))^2 + U(:,i)'*[w4 0; 0 w5]*U(:,i); %cost f 1
    J = J + w1*(X(1,i)-200/3.6)^2 + w2*(X(2,i)-tvxfhg(i,3))^2 + w3*(X(3,i)+X(4,i))^2 + U(:,i)'*[w4 0; 0 w5]*U(:,i); %cost f 2
    % J = J + w1*(X(1,i)-200/3.6)^2 + w3*(X(3,i)^2+X(4,i)^2) + U(:,i)'*[w4 0; 0 w5]*U(:,i); %cost f 3
    % J = J + w1*(X(1,i)-200/3.6)^2 + w3*(X(3,i)^2+X(4,i)^2) + w2*U(1,i)*X(1,i) + U(:,i)'*[w4 0; 0 w5]*U(:,i); %cost f 4 非二阶，引入Fv相乘项
end
%增添一项终端代价
J = J + w6*(X(2,N)-pos_ed_m)^2;
opti.minimize(J);
c_f_end = toc(c_f_start);
fprintf("-目标函数构造耗时：%f s\n",c_f_end);
opti.solver('ipopt'); %内点法
sol = opti.solve();

%% 提取求解结果
t_ = [QsInfo.t_now:QsInfo.step_t:(QsInfo.t_now+QsInfo.set_time+QsInfo.t_err)]'; 
X_opt = sol.value(X);
U_opt = sol.value(U);
v_ms_opt = X_opt(1,:)';  % 速度解
s_opt = X_opt(2,:)';  % 位移解
Ts_opt = X_opt(3,:)'; % 定子温度
Tr_opt = X_opt(4,:)'; % 转子温度

Fe_opt = U_opt(1,:)'; % 牵引力解
Fm_opt = U_opt(2,:)'; % 制动力解

if length(t_) ~= length(tvxfhg)
    if length(t_) > length(tvxfhg)
        t_ = t_(1:length(tvxfhg));
    else
        for i=1:length(tvxfhg)-length(t_)
            t_ = [t_;t_(end)+QsInfo.step_t];
        end
    end
end
res = [t_ v_ms_opt s_opt Ts_opt Tr_opt Fe_opt Fm_opt];

